import os
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel, QPushButton, QHBoxLayout
from PyQt5.QtGui import QPixmap
from PyQt5.QtGui import QImage, QColor
from PyQt5.QtCore import Qt


class SameNameImageViewer(QWidget):
    def __init__(self, parent_folder):
        super().__init__()
        self.parent_folder = parent_folder
        self.image_groups = {}
        self.current_group_index = 0
        self.initUI()

    def initUI(self):
        self.find_same_named_images()
        self.display_images()
        self.create_buttons()

        self.setWindowTitle('同名图片查看器')
        self.show()

    def find_same_named_images(self):

        # list
        folder_names = [
            'val_demo',
            'val_label_demo',
            'U_Net',
            'SegNet',
            'NestedUNet',
            'TVResNet50_PSPNet64',
            'LiteST_Net',
            'LGFO_Net'
        ]

        # 遍历列表中的文件夹名称，依次处理
        for folder_name in folder_names:
            folder_path = os.path.join(self.parent_folder, folder_name, 'predictResult')
            # 检查文件夹是否存在
            if not os.path.exists(folder_path):
                continue
            # 调用处理文件夹的方法
            self.process_folder(folder_path)

        # # 处理 val_demo 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'val_demo/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 val_label_demo 文件夹
        # val_label_demo_folder = os.path.join(self.parent_folder, 'val_label_demo/predictResult')
        # self.process_folder(val_label_demo_folder)
        #
        # # 处理 U_Net1 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'U_Net1/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 SegNet 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'SegNet/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 NestedUNet 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'NestedUNet/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 TVResNet50_PSPNet641 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'TVResNet50_PSPNet641/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 LiteST_Net 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'LiteST_Net/predictResult')
        # self.process_folder(val_demo_folder)
        #
        # # 处理 LGFO_Net 文件夹
        # val_demo_folder = os.path.join(self.parent_folder, 'LGFO_Net/predictResult')
        # self.process_folder(val_demo_folder)



        # 处理其他文件夹
        # for root, dirs, files in os.walk(self.parent_folder):
        #     # if os.path.basename(root) in ['val_demo', 'val_label_demo','U_Net1','SegNet','NestedUNet','TVResNet50_PSPNet641','LiteST_Net','LGFO_Net']:
        #     #     continue
        #     predict_result_path = os.path.join(root, 'predictResult')
        #     if not os.path.exists(predict_result_path):
        #         continue
        #     self.process_folder(predict_result_path)

    def process_folder(self, folder_path):
        for file in os.listdir(folder_path):
            if file.endswith('.png') or file.endswith('.jpg'):
                image_path = os.path.join(folder_path, file)
                image_name = os.path.splitext(file)[0]
                parent_folder = os.path.basename(os.path.dirname(folder_path))
                if image_name not in self.image_groups:
                    self.image_groups[image_name] = []
                self.image_groups[image_name].append((image_path, parent_folder))

    def display_images(self):
        self.layout = QGridLayout(self)
        self.update_display()

    def update_display(self):
        if self.layout is not None:
            for i in reversed(range(self.layout.count())):
                widget = self.layout.itemAt(i).widget()
                if widget is not None:
                    widget.setParent(None)

        row, col = 0, 0
        group_name = list(self.image_groups.keys())[self.current_group_index]
        group_label = QLabel(f"组名: {group_name}")
        self.layout.addWidget(group_label, row, col, 1, 4)
        row += 1
        col = 0
        for image_info in self.image_groups[group_name]:
            image_path, parent_folder = image_info
            image_label = QLabel()
            pixmap = QPixmap(image_path)
            pixmap = pixmap.scaledToWidth(256)  # 根据需要调整宽度，要和图片的宽度一致,否则会出现对比图偏移

            # 获取 val_label_demo 文件夹对应的图像路径
            a_folder_path = os.path.join(self.parent_folder, 'val_label_demo/predictResult')
            a_image_path = os.path.join(a_folder_path, os.path.basename(image_path))

            #设置不需要对比的图片
            if parent_folder in ['val_demo', 'val_label_demo']:
                image_label.setPixmap(pixmap)
            else:
                # 加载图像并进行像素级别的比较
                result_pixmap = self.compare_images(pixmap, QPixmap(a_image_path))
                image_label.setPixmap(result_pixmap)

            # 加载图像并进行像素级别的比较
            # result_pixmap = self.compare_images(pixmap, QPixmap(a_image_path))
            # image_label.setPixmap(result_pixmap)

            image_label.setScaledContents(True)  # 将 pixmap 缩放以适应 QLabel
            self.layout.addWidget(image_label, row, col)
            folder_label = QLabel(parent_folder)
            self.layout.addWidget(folder_label, row + 1, col)
            col = (col + 1) % 4
            if col == 0:
                row += 2

        self.setLayout(self.layout)

    def compare_images(self, pixmap1, pixmap2):
        # 将 QPixmap 转换为 QImage
        image1 = pixmap1.toImage()
        image2 = pixmap2.toImage()
        width = image1.width()
        height = image1.height()
        result_image = QImage(width, height, QImage.Format_ARGB32)

        # 遍历像素进行比较
        for y in range(height):
            for x in range(width):
                pixel1 = QColor(image1.pixel(x, y))
                pixel2 = QColor(image2.pixel(x, y))


                # 如果像素不同，根据像素的颜色分别判断 TP、FP、TN、FN，并设置相应的颜色
                if pixel1 == QColor(Qt.white) and pixel2 == QColor(Qt.black):
                    result_image.setPixelColor(x, y, QColor(Qt.green))  # FN
                elif pixel1 == QColor(Qt.black) and pixel2 == QColor(Qt.white):
                    result_image.setPixelColor(x, y, QColor(Qt.red))  # FP
                elif pixel1 == QColor(Qt.black) and pixel2 == QColor(Qt.black):
                    result_image.setPixelColor(x, y, QColor(Qt.black))  # TN
                elif pixel1 == QColor(Qt.white) and pixel2 == QColor(Qt.white):
                    result_image.setPixelColor(x, y, QColor(Qt.white))  # TP

        return QPixmap.fromImage(result_image)

    def create_buttons(self):
        button_layout = QHBoxLayout()

        prev_button = QPushButton("上一组")
        prev_button.clicked.connect(self.previous_group)
        button_layout.addWidget(prev_button)

        next_button = QPushButton("下一组")
        next_button.clicked.connect(self.next_group)
        button_layout.addWidget(next_button)

        self.layout.addLayout(button_layout, len(self.image_groups) * 2 + 1, 0, 1, 4)

    def previous_group(self):
        if self.current_group_index > 0:
            self.current_group_index -= 1
            self.update_display()

    def next_group(self):
        if self.current_group_index < len(self.image_groups) - 1:
            self.current_group_index += 1
            self.update_display()

if __name__ == '__main__':
    app = QApplication([])
    parent_folder = '/home/jolly/develop/PycharmProjects/remote_sensor/bijie256'  # 替换为您的父文件夹路径
    window = SameNameImageViewer(parent_folder)
    app.exec_()
